from typing import Dict, Any
import os
from pathlib import Path

class TranslationConfig:
    """
    翻译引擎配置类
    支持从环境变量和配置文件读取配置
    """
    
    # 默认配置
    DEFAULT_CONFIG = {
        # 引擎类型选择
        "engine_type": "whisper_marian",  # 默认使用分步处理
        
        # WhisperMarian 配置
        "whisper_marian": {
            "asr_model": "base",
            "translation_model": "Helsinki-NLP/opus-mt-zh-en"
        },
        
        # SeamlessM4T 配置
        "seamless_m4t": {
            "model_name": "facebook/seamless-m4t-medium",
            "device": "auto"  # auto, cpu, cuda
        },
        
        # Triton 配置（符合规范要求）
        "triton": {
            "server_url": "localhost:8001",
            "model_name": "whisper_asr",
            "protocol": "grpc",  # grpc 或 http
            "use_ssl": False,
            "timeout": 30.0
        },
        
        # 通用配置
        "common": {
            "max_audio_duration": 300,  # 最大音频时长（秒）
            "supported_languages": ["zh", "en", "ja", "ko", "fr", "de", "es", "ru"],
            "confidence_threshold": 0.7,  # 置信度阈值
            "enable_segments": True,  # 是否启用分段
            "cache_results": True,  # 是否缓存结果
            "log_level": "INFO"
        }
    }
    
    def __init__(self, config_file: str = None):
        """
        初始化配置
        
        Args:
            config_file: 配置文件路径（可选）
        """
        self.config = self.DEFAULT_CONFIG.copy()
        
        # 从环境变量加载配置
        self._load_from_env()
        
        # 从配置文件加载配置
        if config_file and os.path.exists(config_file):
            self._load_from_file(config_file)
    
    def _load_from_env(self):
        """从环境变量加载配置"""
        env_mapping = {
            "TRANSLATION_ENGINE_TYPE": "engine_type",
            "TRITON_SERVER_URL": ("triton", "server_url"),
            "TRITON_MODEL_NAME": ("triton", "model_name"),
            "TRITON_PROTOCOL": ("triton", "protocol"),
            "SEAMLESS_M4T_MODEL": ("seamless_m4t", "model_name"),
            "CONFIDENCE_THRESHOLD": ("common", "confidence_threshold"),
            "MAX_AUDIO_DURATION": ("common", "max_audio_duration")
        }
        
        for env_var, config_path in env_mapping.items():
            value = os.getenv(env_var)
            if value is not None:
                if isinstance(config_path, str):
                    self.config[config_path] = value
                else:
                    section, key = config_path
                    if section not in self.config:
                        self.config[section] = {}
                    self.config[section][key] = value
    
    def _load_from_file(self, config_file: str):
        """从配置文件加载配置"""
        try:
            import yaml
            with open(config_file, 'r', encoding='utf-8') as f:
                file_config = yaml.safe_load(f)
                self._merge_config(file_config)
        except ImportError:
            # 如果没有yaml，尝试JSON
            try:
                import json
                with open(config_file, 'r', encoding='utf-8') as f:
                    file_config = json.load(f)
                    self._merge_config(file_config)
            except Exception as e:
                print(f"加载配置文件失败: {e}")
    
    def _merge_config(self, new_config: Dict):
        """合并配置"""
        for key, value in new_config.items():
            if key in self.config and isinstance(self.config[key], dict) and isinstance(value, dict):
                self.config[key].update(value)
            else:
                self.config[key] = value
    
    def get_engine_type(self) -> str:
        """获取当前选择的引擎类型"""
        return self.config.get("engine_type", "whisper_marian")
    
    def get_engine_config(self, engine_type: str = None) -> Dict[str, Any]:
        """获取指定引擎的配置"""
        if engine_type is None:
            engine_type = self.get_engine_type()
        
        return self.config.get(engine_type, {})
    
    def get_common_config(self) -> Dict[str, Any]:
        """获取通用配置"""
        return self.config.get("common", {})
    
    def set_engine_type(self, engine_type: str):
        """设置引擎类型"""
        if engine_type in ["whisper_marian", "seamless_m4t", "triton"]:
            self.config["engine_type"] = engine_type
        else:
            raise ValueError(f"不支持的引擎类型: {engine_type}")
    
    def update_engine_config(self, engine_type: str, config: Dict[str, Any]):
        """更新指定引擎的配置"""
        if engine_type not in self.config:
            self.config[engine_type] = {}
        self.config[engine_type].update(config)
    
    def validate_config(self) -> bool:
        """验证配置是否有效"""
        engine_type = self.get_engine_type()
        
        # 检查引擎配置是否存在
        if engine_type not in self.config:
            return False
        
        # 检查必需配置
        if engine_type == "triton":
            triton_config = self.config["triton"]
            required_keys = ["server_url", "model_name"]
            if not all(key in triton_config for key in required_keys):
                return False
        
        return True
    
    def get_config_summary(self) -> Dict[str, Any]:
        """获取配置摘要"""
        return {
            "current_engine": self.get_engine_type(),
            "available_engines": ["whisper_marian", "seamless_m4t", "triton"],
            "engine_config": self.get_engine_config(),
            "common_config": self.get_common_config(),
            "is_valid": self.validate_config()
        }
    
    def save_config(self, config_file: str):
        """保存配置到文件"""
        try:
            import yaml
            with open(config_file, 'w', encoding='utf-8') as f:
                yaml.dump(self.config, f, default_flow_style=False, allow_unicode=True)
        except ImportError:
            # 如果没有yaml，使用JSON
            import json
            with open(config_file, 'w', encoding='utf-8') as f:
                json.dump(self.config, f, ensure_ascii=False, indent=2)

# 全局配置实例
translation_config = TranslationConfig()

# 便捷函数
def get_translation_config() -> TranslationConfig:
    """获取全局翻译配置"""
    return translation_config

def get_current_engine_type() -> str:
    """获取当前引擎类型"""
    return translation_config.get_engine_type()

def get_engine_config(engine_type: str = None) -> Dict[str, Any]:
    """获取引擎配置"""
    return translation_config.get_engine_config(engine_type)




